import jwt
import base64
from typing import Optional, Dict, Any
from datetime import datetime, timedelta, timezone
from biz.utils.logger import logger

class JwtUtil:
    """JWT工具类 - Python版本"""
    
    JWT_SECRET_KEY = "cOM0987tElCHiNa/6543JwTCeNtER21=="
    
    @classmethod
    def get_encode_key(cls) -> str:
        """获取编码后的密钥"""
        encoded_key = base64.b64encode(cls.JWT_SECRET_KEY.encode()).decode()
        return encoded_key
    
    @classmethod
    def parse_jwt(cls, jwt_token: str) -> Optional[Dict[str, Any]]:
        """解析JWT token
        
        Args:
            jwt_token: JWT token字符串
            
        Returns:
            解析成功返回claims字典，失败返回None
        """
        try:
            # 使用base64编码的密钥解析JWT
            secret_key = base64.b64decode(cls.get_encode_key())
            claims = jwt.decode(jwt_token, secret_key, algorithms=['HS256'])
            return claims
        except jwt.ExpiredSignatureError:
            logger.info("签名过期")
            return None
        except jwt.InvalidSignatureError:
            logger.info("签名错误")
            return None
        except jwt.InvalidTokenError:
            logger.info("无效的token")
            return None
        except Exception as e:
            logger.error(f"解析JWT时发生未知错误: {str(e)}")
            return None
    
    @classmethod
    def generate_jwt(cls, payload: Dict[str, Any], expires_in_hours: int = 24) -> str:
        """生成JWT token
        
        Args:
            payload: 要编码的数据
            expires_in_hours: 过期时间（小时）
            
        Returns:
            生成的JWT token字符串
        """
        try:
            # 添加过期时间
            payload['exp'] = datetime.now(timezone.utc) + timedelta(hours=expires_in_hours)
            payload['iat'] = datetime.now(timezone.utc)
            
            # 使用base64编码的密钥生成JWT
            secret_key = base64.b64decode(cls.get_encode_key())
            token = jwt.encode(payload, secret_key, algorithm='HS256')
            return token
        except Exception as e:
            logger.error(f"生成JWT时发生错误: {str(e)}")
            raise
    
    @classmethod
    def is_token_valid(cls, jwt_token: str) -> bool:
        """检查token是否有效
        
        Args:
            jwt_token: JWT token字符串
            
        Returns:
            token有效返回True，否则返回False
        """
        return cls.parse_jwt(jwt_token) is not None
    
    @classmethod
    def get_user_id_from_token(cls, jwt_token: str) -> Optional[str]:
        """从token中获取用户ID
        
        Args:
            jwt_token: JWT token字符串
            
        Returns:
            用户ID字符串，获取失败返回None
        """
        claims = cls.parse_jwt(jwt_token)
        if claims:
            return claims.get('user_id') or claims.get('sub')
        return None


if __name__ == '__main__':
    # 解析JWT token
    token = "your_jwt_token_here"
    claims = JwtUtil.parse_jwt(token)
    if claims:
        print(f"Token有效，用户信息: {claims}")
    else:
        print("Token无效或已过期")

    # 生成新的JWT token
    payload = {"user_id": "123", "username": "admin"}
    new_token = JwtUtil.generate_jwt(payload, expires_in_hours=24)
    print(f"生成的token: {new_token}")

    # 验证token有效性
    if JwtUtil.is_token_valid(token):
        user_id = JwtUtil.get_user_id_from_token(token)
        print(f"用户ID: {user_id}")
